# -*- coding: utf-8 -*-
import time
import subprocess
import os
import signal
import sys
import random
from tuner import JvmAutoTuner
import optuna
from optuna.pruners import MedianPruner

# --- 模拟的应用实例句柄 ---
class AppInstance:
    def __init__(self, pid: int, jvm_opts: str, port: int):
        self.pid = pid
        self.jvm_opts = jvm_opts
        self.port = port
        self.url = f"http://localhost:{port}"

    def __repr__(self):
        return f"AppInstance(pid={self.pid}, port={self.port})"

# --- 用户需要自定义的回调函数 ---

def my_launch_callback(jvm_opts: str) -> AppInstance:
    """【用户实现】启动你的Java应用。"""
    port = 8080
    command = [sys.executable, "-m", "http.server", str(port)]
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, preexec_fn=os.setsid)
    time.sleep(1)
    if process.poll() is not None:
        raise RuntimeError("Application failed to start.")
    return AppInstance(pid=process.pid, jvm_opts=jvm_opts, port=port)


def my_evaluate_callback(instance: AppInstance, trial: optuna.trial.Trial) -> dict:
    """【用户实现】评估应用性能并返回一个包含多项指标的字典。"""
    try:
        total_duration = 15
        step_duration = 5
        for step in range(1, (total_duration // step_duration) + 1):
            current_time = step * step_duration
            print(f"  [Evaluator] Running load test... ({current_time}/{total_duration}s)")
            time.sleep(step_duration)
            intermediate_metrics = calculate_mock_metrics(instance.jvm_opts, current_time / total_duration)
            trial.report(intermediate_metrics['latency'], step=step)
            if trial.should_prune():
                raise optuna.exceptions.TrialPruned()
        final_metrics = calculate_mock_metrics(instance.jvm_opts, 1.0)
        return final_metrics
    finally:
        try:
            os.killpg(os.getpgid(instance.pid), signal.SIGTERM)
        except ProcessLookupError:
            pass

def calculate_mock_metrics(jvm_opts: str, completion_ratio: float) -> dict:
    """【模拟核心】根据JVM参数和测试完成度计算一个模拟性能指标字典。"""
    base_latency = 300
    base_throughput = 2000
    if "-XX:+UseG1GC" in jvm_opts:
        base_latency *= 0.6
        base_throughput *= 1.5
    elif "-XX:+UseZGC" in jvm_opts:
        base_latency *= 0.5
        base_throughput *= 1.3
    elif "-XX:+UseParallelGC" in jvm_opts:
        base_latency *= 1.2
        base_throughput *= 0.9
    mem_mb = 4096
    xmx_opt = [s for s in jvm_opts.split() if s.startswith("-Xmx")]
    if xmx_opt:
        mem_mb = int(xmx_opt[0].replace("-Xmx", "").replace("m", ""))
    mem_penalty = 1 + abs(mem_mb - 6144) / 8192 # 假设最佳内存在6G左右
    final_latency = base_latency * mem_penalty
    final_throughput = base_throughput / mem_penalty
    noise = random.uniform(0.9, 1.1)
    stability_factor = 0.8 + 0.2 * completion_ratio
    return {
        'latency': final_latency * noise / stability_factor,
        'throughput': final_throughput * noise * stability_factor
    }

# --- 主程序入口 ---
if __name__ == "__main__":
    # 1. 定义超大规模、高度详细的参数搜索空间
    param_space = {
        # --- 1. 核心GC算法选择 ---
        "GC": {"type": "categorical", "choices": ["G1GC", "ZGC", "ParallelGC"]},

        # --- 2. 堆内存核心参数 ---
        "Xmx": {"type": "int", "low": 4096, "high": 12288, "step": 2048, "unit": "m"},
        # Xms将通过框架逻辑自动设置为与Xmx相等，这里只需定义即可
        "Xms": {"type": "int", "low": 4096, "high": 12288, "is_constraint_target": True},
        "NewRatio": {"type": "int", "low": 2, "high": 4},
        "SurvivorRatio": {"type": "int", "low": 6, "high": 10},

        # --- 3. G1GC 专用参数 ---
        "MaxGCPauseMillis": {
            "type": "int", "low": 100, "high": 300, "step": 50,
            "condition": ("GC", "G1GC")
        },
        "G1HeapRegionSize": {
            "type": "categorical", "choices": [8, 16, 32], "unit": "m",
            "condition": ("GC", "G1GC")
        },
        "InitiatingHeapOccupancyPercent": {
            "type": "int", "low": 35, "high": 55, "step": 5,
            "condition": ("GC", "G1GC")
        },

        # --- 4. ParallelGC 专用参数 ---
        "ParallelGCThreads": {
            "type": "int", "low": 4, "high": 16, "step": 4,
            "condition": ("GC", "ParallelGC")
        },
        "UseAdaptiveSizePolicy": {
            "type": "categorical", "choices": [True, False], "is_flag": True,
            "condition": ("GC", "ParallelGC")
        },

        # --- 5. JIT编译器与代码缓存 ---
        "TieredCompilation": {"type": "categorical", "choices": [True, False], "is_flag": True},
        "CompileThreshold": {"type": "int", "low": 5000, "high": 15000, "step": 2500},

        # --- 6. 内存管理与性能优化 ---
        "AlwaysPreTouch": {"type": "categorical", "choices": [True, False], "is_flag": True},
        "UseLargePages": {"type": "categorical", "choices": [True, False], "is_flag": True},
        "UseBiasedLocking": {"type": "categorical", "choices": [True, False], "is_flag": True},
        "UseNUMA": {"type": "categorical", "choices": [True, False], "is_flag": True},
        "MaxTenuringThreshold": {"type": "int", "low": 4, "high": 15},
        "ObjectAlignmentInBytes": {"type": "categorical", "choices": [8, 16, 32]}
    }

    # 2. 实例化并运行调优器
    tuner = JvmAutoTuner(
        param_space=param_space,
        launch_callback=my_launch_callback,
        evaluate_callback=my_evaluate_callback,
        goal="balanced",
        n_trials=50,
        timeout=600,
        pruner=MedianPruner(n_startup_trials=5, n_warmup_steps=1)
    )
    
    best_params, best_score = tuner.run()

    # 3. 打印最终结果
    if best_params:
        print("\n--- Optimal Configuration Found ---")
        print(f"Best Score: {best_score:.4f}")
        print("Best Parameters:")
        for param, value in best_params.items():
            print(f"  - {param}: {value}")

